'''
    1. 将深度图00xx_depth.tiff转化为点云,并使用mayavi显示  v
    2. 画3D bbox
'''

import mayavi.mlab as mlab
import cv2
import sys
sys.path.append('./')
sys.path.append('./pose_estimator')
from pose_estimator.lib.utils import load_depth
import numpy as np
import _pickle as cPickle

data_dir = 'rs_imgs/'
img_width = 640
img_height = 480
depth_scale = 1000


def get_3d_bbox(size, shift=0):
    """
    Args:
        size: [3] or scalar
        shift: [3] or scalar
    Returns:
        bbox_3d: [3, N]
    """
    bbox_3d = np.array([[+size[0] / 2, +size[1] / 2, +size[2] / 2],
                        [+size[0] / 2, +size[1] / 2, -size[2] / 2],
                        [-size[0] / 2, +size[1] / 2, +size[2] / 2],
                        [-size[0] / 2, +size[1] / 2, -size[2] / 2],
                        [+size[0] / 2, -size[1] / 2, +size[2] / 2],
                        [+size[0] / 2, -size[1] / 2, -size[2] / 2],
                        [-size[0] / 2, -size[1] / 2, +size[2] / 2],
                        [-size[0] / 2, -size[1] / 2, -size[2] / 2]]) + shift
    bbox_3d = bbox_3d.transpose()
    return bbox_3d

def transform_coordinates_3d(coordinates, sRT):
    """
    Args:
        coordinates: [3, N]
        sRT: [4, 4]
    Returns:
        new_coordinates: [3, N]
    """
    assert coordinates.shape[0] == 3
    coordinates = np.vstack([coordinates, np.ones((1, coordinates.shape[1]), dtype=np.float32)])
    new_coordinates = sRT @ coordinates
    new_coordinates = new_coordinates[:3, :] / new_coordinates[3, :]
    #new_coordinates = new_coordinates[:3, :]
    return new_coordinates

def draw_detections(pred_sRT, pred_size):
    """ Visualize pose predictions.
    """
    transformed_bbox_3d = np.zeros((pred_sRT.shape[0], 3, 9))
    for i in range(pred_sRT.shape[0]):
        sRT = pred_sRT[i, :, :]
        bbox_3d = get_3d_bbox(pred_size[i, :], 0)
        transformed_bbox_3d[i, :, :8] = transform_coordinates_3d(bbox_3d, sRT)
        transformed_bbox_3d[i, :, 8] = np.mean(transformed_bbox_3d[i, :, :8], axis=1)
    return transformed_bbox_3d


cam_fx, cam_fy, cam_cx, cam_cy = 603.1332397460938, 601.6819458007812, 323.1732177734375, 241.5373077392578
intrinsics = [[cam_fx, 0, cam_cx], [0, cam_fy, cam_cy], [0, 0, 1]]
intrisic_inv = np.linalg.inv(intrinsics)


img_id_list = []
for i in range(0,17):
    img_id_list.append(str(i).zfill(4))
for img_id in img_id_list:
    pts = np.zeros((img_width * img_height, 3))    
    depth_img = load_depth(data_dir + img_id) # shape:(540, 960)
    #rgb_img = cv2.imread(data_dir + img_id + '_color.png')
    #depth_img = depth_img / depth_scale
    uv_list = np.meshgrid(range(0, img_width), range(0, img_height))
    u = uv_list[0][np.newaxis, :]
    v = uv_list[1][np.newaxis, :]
    z = depth_img[np.newaxis, :]
    zzz = np.concatenate((z, z, z), 0)  # (3, 540, 960)
    uv1 = np.concatenate((u, v, np.ones_like(u)), 0)  # (3, 540, 960)
    uv1 = uv1.reshape(3, -1)  # (3, 540 x 960)
    temp =intrisic_inv @ uv1
    XcYcZc = np.multiply(temp, zzz.reshape(3, -1))
    #print(XcYcZc)

    #mlab.plot3d(XcYcZc[0, :], XcYcZc[1, :], XcYcZc[2, :], lambda xc, yc, zc : xc+yc+zc, colormap="autumn")
    # 画3D bbox
    with open(data_dir + img_id + '_pose.pkl', 'rb') as f:
        result = cPickle.load(f)
    mugs_3Dbbox = draw_detections(result['pred_RTs'], result['pred_scales']) * depth_scale
    mlab.points3d(mugs_3Dbbox[:, 0, :].reshape(-1), mugs_3Dbbox[:, 1, :].reshape(-1), mugs_3Dbbox[:, 2, :].reshape(-1), mugs_3Dbbox[:, 2, :].reshape(-1), scale_factor=10.0, scale_mode="none", colormap="spectral")
    #mlab.plot3d(mugs_3Dbbox[:, 0, :].reshape(-1), mugs_3Dbbox[:, 1, :].reshape(-1), mugs_3Dbbox[:, 2, :].reshape(-1), mugs_3Dbbox[:, 2, :].reshape(-1), colormap="spectral", line_width=10000000000.0)
    mlab.points3d(XcYcZc[0, :], XcYcZc[1, :], XcYcZc[2, :], lambda xc, yc, zc : xc+yc+zc, colormap="autumn", scale_mode="none", scale_factor=3)
    mlab.show()